import os
import cv2
import glob
import argparse
import time

color=[(0,0,255),
       (255,0,0),
       (0,255,0),
       (255,255,0),
       (255,0,255),
       (0,255,255),
       (255,255,255),
       (128,128,0),
       (128,0,128),
       (0,128,128),
       (128,128,128),]


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--path',help='json path')
    parser.add_argument('--outpath',help='json path')
    args = parser.parse_args()
    return args

def draw(path,outpath):
    if not  os.path.exists(outpath):
        os.makedirs(outpath)
    txtlist= glob.glob(path+'/*.txt')
    for txt in txtlist:
        labelfile = txt
        imagefile = txt[:-3]+'jpg'
        img = cv2.imread(imagefile)
        file = open(labelfile,'r')
        height,width,c = img.shape    
        for line in file:
            c,cx,cy,w,h = line.split(' ')
            cx,cy,w,h = float(cx),float(cy),float(w),float(h)
            x1,y1,x2,y2 = int((cx-w/2)*width),int((cy-h/2)*height),int((cx+w/2)*width),int((cy+h/2)*height)
            cv2.rectangle(img, (x1,y1), (x2,y2), color[int(c)], 2)
        cv2.imwrite(outpath+ '/'+str(time.time())+'.jpg', img)
        print(txt)

if __name__=='__main__':
    args = parse_args()
    draw(args.path,args.outpath)
